[Feature][MRV2] Adapt PCP for data parallelism by wzx0726 · Pull Request #15969 · vllm-project/vllm-ascend

wzx0726 · 2026-09-07T14:08:44Z

What this PR does / why we need it?

Add Ascend MRV2 handling for PCP combined with DP, reusing the upstream PCP dispatch sizing and input partitioning paths.

When a DP replica becomes idle, its dummy batch bypasses PCP partitioning. Saved PCP attention state may be absent or belong to the preceding real batch. Runtime dummy inputs use runner buffers, while decode graphs capture persistent PCP buffers.

Build dummy PCP attention contexts from the current batch, block tables and rank-local slot mappings, including the DSA metadata path.
Refresh persistent PCP buffers before dummy execution. Keep the implementation in AscendPCPManager; NPUModelRunner.prepare_dummy_attn delegates or takes the existing non-PCP path.
Synchronize replicated speculative drafts independently of the target's PCP-local DP state, with every DP replica participating, including idle and decode replicas.
Restrict PCP+DP graph configuration to eager or FULL_DECODE_ONLY.

Upstream validation dependency: vllm-project/vllm#54523, merged as 7c2f1ff4958eaf0818405e9192c71608fe4a16b1, moves the PCP+DP restriction from common ParallelConfig validation to CUDA/ROCm platform checks. This PR relies on that change and contains no Ascend validator bypass or Pydantic schema rebuilding.

Related upstream PR: vllm-project/vllm#53867 ([Feature][PCP] Support decode-only FULL CUDA graphs).

Once the paired vLLM includes #53867, adapt its PCP prepare_inputs_to_capture entry point to create AscendInputBatch directly in persistent PCP buffers. After capture and real-to-idle replay validation, remove AscendPCPManager.prepare_dummy_attn and the runner override, as tracked by TODO(wzx0726). The Ascend dummy attention-context handling and replicated-draft DP synchronization remain necessary.

Does this PR introduce any user-facing change?

Adds Ascend-side MRV2 handling for PCP combined with DP. No new CLI options or environment variables are introduced.

Rebased onto Ascend main ad86348b0cb2324d643df6a496f2d9c3879e4481. The draft conflict resolution preserves vLLM 0.28.0 num_tokens_across_dp forwarding and applies fresh synchronization for replicated drafts on the main2main dp_sync interface.

Compatibility and remaining validation:

The previously tested baseline was Ascend b1c91857c16bde10c2fa7f5d7548b7666a86d4bb with vLLM e6bfe03ad73a3330cb427885aa90d97a12e1c704.
Removing the validator bypass requires vLLM #54523. Ascend main currently pins b2f685834a6456197e7033966fdef52a23f1abcd, which predates that change and still rejects PCP+DP. The local vLLM checkout remains at the user-selected 7c2f1ff4958eaf0818405e9192c71608fe4a16b1. This PR does not change the main2main pin or reintroduce the bypass.
Compatibility work against that newer revision is still pending, including the added prepare_dummy_attn argument and removal of InputBatch.max_seq_len_np. The current branch is not yet validated to run against that revision.
Multi-card Attention/MLA/DSA inference, real-to-idle DP transitions during FULL_DECODE_ONLY replay, speculative-decoding numerical correctness, and EP/TP combinations require model-level validation on the final paired sources.

How was this patch tested?

The latest rebase onto ad86348b0 resolves an append-location conflict in test_pcp_manager_v2.py, retaining both the main DCP padding regression and this PR's tests. AST comparison confirms that all main tests and the final PR tests are preserved, and the other changed files match the automatic merge. Ruff lint/format, syntax parsing, and git diff --check passed. Runtime tests were not rerun for this conflict-only update.

For the preceding rebase onto f8481287a, 6 targeted candidate-method cases passed in an isolated process in the Ascend container: five main2main draft-sync cases and one mocked v0.28.0 argument-routing case. The exact candidate propose method was loaded into the installed runtime classes; native dependencies and the rest of the runtime remained at their installed versions. This is method-level regression evidence, not full rebased-source or dual-version runtime validation.

All eight changed Python files passed Ruff lint/format checks, syntax parsing, and git diff --check. Full unit-suite and model validation on the final paired sources remain pending.

The runtime adaptation suite previously passed on the earlier baseline in an isolated Ascend-container test directory, using CPU tensors and mocks for device kernels:

python3 -m pytest -q tests/ut/worker/test_pcp_manager_v2.py tests/ut/worker/test_model_runner_v2.py tests/ut/worker/test_attn_utils_v2.py tests/ut/worker/test_mtp_pcp_speculator_v2.py

75 passed (14 existing torch.jit deprecation warnings). This is prior-baseline evidence, not a test result for vLLM 7c2f1ff.
Covers missing/stale dummy PCP state, rank-local mappings, persistent buffer contents and storage, DSA metadata forwarding, non-PCP fallback, graph restrictions, and replicated draft DP synchronization.
The removed validator bypass's dedicated tests are also removed; the existing unrelated platform test is retained.
For the removal, targeted Ruff lint/format checks, Python syntax checks and git diff --check passed. No runtime tests were rerun against the upgraded vLLM. The local format.sh ci entry point remains unavailable because its shell lacks pre-commit.
vLLM main: vllm-project/vllm@b2f6858

gemini-code-assist · 2026-09-07T14:09:10Z

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request adapts the Ascend MRV2 PCP (Prefill Context Parallelism) to support data parallelism. It addresses discrepancies between DP metadata and partitioned model inputs by calculating padded execution lengths prior to dispatch and ensuring that input preparation correctly restores token totals. Additionally, it improves robustness by enabling the construction of dummy attention contexts for runtime batches that bypass standard partitioning.

Highlights

PCP and Data Parallelism Integration: Implemented logic to compute PCP-padded execution lengths before DP dispatch, ensuring DP metadata remains consistent with partitioned model inputs.
Dummy Attention Context Construction: Added support for building runtime dummy attention contexts from current dummy batches, block tables, and slot mappings to handle cases where saved PCP state is missing or stale.
Configuration Constraints: Restricted PCP+DP graph configurations to eager mode or FULL_DECODE_ONLY to ensure stability.

New Features

🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request implements support for Prefill Context Parallelism (PCP) combined with Data Parallelism (DP) in Model Runner V2, including dispatch size calculations, config validation, and dummy batch handling for attention context. Feedback highlights two critical issues: first, a potential race condition/stale data issue in gather_batch_req_state where num_computed_tokens_np is read before synchronization when speculative decoding is active; second, a potential RuntimeError in build_attention_context due to calling .view() on a potentially non-contiguous slot_mappings tensor, which should be replaced with .reshape().\n\nSuggested PR Title:\n\nmarkdown\n[Ops][Feature] Support PCP+DP in Model Runner V2 and handle dummy batches\n\n\nSuggested PR Summary:\n\nmarkdown\n### What this PR does / why we need it?\nThis PR introduces support for Prefill Context Parallelism (PCP) combined with Data Parallelism (DP) in Model Runner V2. It adds a mechanism to query PCP execution length for dispatch, updates input preparation to handle PCP-local dispatch totals during prefill, and enhances `build_attention_context` to support runtime dummy batches. Additionally, it restricts PCP+DP to eager mode or FULL_DECODE_ONLY CUDA graphs.\n\n### Does this PR introduce _any_ user-facing change?\nNo.\n\n### How was this patch tested?\nTested via new unit tests in `test_attn_utils_v2.py`, `test_model_runner_v2.py`, and `test_pcp_manager_v2.py`.\n

gemini-code-assist · 2026-09-07T14:10:39Z

+            if (
+                self.pcp_manager is not None
+                and self.dp_size > 1
+                and batch_state is not None
+                and batch_state.has_prefill
+            ):
+                assert isinstance(self.pcp_manager, AscendPCPManager)
+                # Dispatch needs the PCP-padded execution total; per-request
+                # counts remain global and prepare_inputs restores their sum.
+                num_tokens = self.pcp_manager.get_num_tokens_for_dispatch(
+                    batch_state.num_scheduled_tokens,
+                    self.req_states.num_computed_tokens_np[batch_state.idx_mapping_np],
+                    batch_state.is_prefilling_np,
+                )
+                batch_state = batch_state._replace(num_tokens=num_tokens)


When speculative decoding is enabled (self.speculator is not None), the GPU-side num_computed_tokens may have been reverted/corrected due to token rejections. This corrected state is copied asynchronously to self.num_computed_tokens_cpu at the end of the previous step, but it is not synchronized and copied to self.req_states.num_computed_tokens_cpu (which updates num_computed_tokens_np) until _update_seq_lens_cpu is called inside prepare_inputs. Since gather_batch_req_state is called before prepare_inputs, reading self.req_states.num_computed_tokens_np here will use stale computed token counts, leading to an incorrect PCP dispatch size calculation. We should synchronize and update the CPU computed tokens for the current batch's requests before calling get_num_tokens_for_dispatch.

if ( self.pcp_manager is not None and self.dp_size > 1 and batch_state is not None and batch_state.has_prefill ): assert isinstance(self.pcp_manager, AscendPCPManager) if self.speculator is not None: self.num_computed_tokens_event.synchronize() for req_id in batch_state.req_ids: req_index = self.req_states.req_id_to_index[req_id] self.req_states.num_computed_tokens_cpu[req_index] = self.num_computed_tokens_cpu[req_index] # Dispatch needs the PCP-padded execution total; per-request # counts remain global and prepare_inputs restores their sum. num_tokens = self.pcp_manager.get_num_tokens_for_dispatch( batch_state.num_scheduled_tokens, self.req_states.num_computed_tokens_np[batch_state.idx_mapping_np], batch_state.is_prefilling_np, ) batch_state = batch_state._replace(num_tokens=num_tokens)

gemini-code-assist · 2026-09-07T14:10:39Z

+                global_slot_mappings=slot_mappings.view(slot_mappings.shape[0], self.pcp_world_size, num_tokens)[
+                    :, self.pcp_rank
+                ],


Using .view() on slot_mappings can raise a RuntimeError if the tensor is non-contiguous (e.g., if it was created via slicing or other non-contiguous operations). Since slot_mappings is passed from external callers and its contiguity is not guaranteed, it is safer to use .reshape() instead of .view().

Suggested change

global_slot_mappings=slot_mappings.view(slot_mappings.shape[0], self.pcp_world_size, num_tokens)[

:, self.pcp_rank

],

global_slot_mappings=slot_mappings.reshape(slot_mappings.shape[0], self.pcp_world_size, num_tokens)[

:, self.pcp_rank

],

github-actions · 2026-09-07T14:11:16Z